/*
    Copyright 2007-2011 Patrik Jonsson, sunrise@familjenjonsson.org. 
    Based on code by Volker Springel, used with permission.

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/** \file

    Implementation of the cooling class, which calculates cooling
    rates and times.
*/

// $Id$

#include <cmath>
#include "cooling.h"
#include <cstdio>
#include <stdexcept>

using namespace std;

double cooling::CoolingTime(double u, double rho, double ne_guess) const
{
  // hydrogen number density
  const double nH = x_H * rho / m_p;
  const double ratefact = nH * nH / rho;

  const double temp = ConvertUToTemp(u, rho, ne_guess);
  
  const double LambdaNet = CoolingRate(temp, rho, ne_guess);

  const double coolingtime = u / (ratefact * LambdaNet);

  return coolingtime;
}

double cooling::CoolingRate(double T, double rho, double ne_guess) const
{
  const pair<cooling_abundances, cooling_rates> va = 
    ComputeAbundancesAndRates(T, rho, ne_guess);
  const cooling_rates& v = va.second;
  const cooling_abundances& a = va.first;

  // Compute cooling and heating rate (cf KWH Table 1) in units of nH**2 

  // excitation
  const double LambdaExcH0 = v.BetaH0 * a.ne * a.nH0;
  const double LambdaExcHep = v.BetaHep * a.ne * a.nHep;
  const double LambdaExc = LambdaExcH0 + LambdaExcHep;

  // ionization
  const double LambdaIonH0 = 2.18e-11 * v.GammaeH0 * a.ne * a.nH0;
  const double LambdaIonHe0 = 3.94e-11 * v.GammaeHe0 * a.ne * a.nHe0;
  const double LambdaIonHep = 8.72e-11 * v.GammaeHep * a.ne * a.nHep;
  const double LambdaIon = LambdaIonH0 + LambdaIonHe0 + LambdaIonHep;

  // recombination
  const double LambdaRecHp = 1.036e-16 * T * a.ne * (v.AlphaHp * a.nHp);
  const double LambdaRecHep = 1.036e-16 * T * a.ne * (v.AlphaHep * a.nHep);
  const double LambdaRecHepp = 1.036e-16 * T * a.ne * (v.AlphaHepp * a.nHepp);
  const double LambdaRecHepd = 6.526e-11 * v.Alphad * a.ne * a.nHep;
  const double LambdaRec = LambdaRecHp + LambdaRecHep + LambdaRecHepp +
    LambdaRecHepd;
      
  // free-free
  const double LambdaFF = v.Betaff * (a.nHp + a.nHep + 4 * a.nHepp) * a.ne;

  const double Lambda = LambdaExc + LambdaIon + LambdaRec + LambdaFF;

  return Lambda;
}


double cooling::ConvertUToTemp(double u, double rho, double& ne_guess) const
{
  double temp, temp_old, temp_new, max = 0, ne_old;
  double mu;
  int iter = 0;

  mu = (1 + 4 * yhelium) / (1 + yhelium + ne_guess);
  temp = (gamma - 1) / k * u * m_p * mu;


  do {
    ne_old = ne_guess;

    const pair<cooling_abundances, cooling_rates> anr = 
      ComputeAbundancesAndRates(temp, rho, ne_guess);
    const cooling_abundances& a(anr.first);
    temp_old = temp;
    ne_guess = a.ne;

    mu = (1 + 4 * yhelium) / (1 + yhelium + a.ne);

    temp_new = (gamma - 1) / k * u * m_p * mu;

    max =
      std::max(max,
	       temp_new / (1 + yhelium + a.ne) * 
	       fabs((a.ne - ne_old) / (temp_new - temp_old + 1.0)));

    temp = temp_old + (temp_new - temp_old) / (1 + max);
    iter++;

    if(iter > (maxiter - 10))
      printf("-> temp= %g ne=%g\n", temp, a.ne);
  }
  while(fabs(temp - temp_old) > 1.0e-3 * temp && iter < maxiter);

  if(iter >= maxiter) {
    printf("failed to converge in convert_u_to_temp()\n");
    //printf("u_input= %g\nrho_input=%g\n ne_input=%g\n", u_input, rho_input, ne_input);
    //printf("DoCool_u_old_input=%g\nDoCool_rho_input= %g\nDoCool_dt_input= %g\nDoCool_ne_guess_input= %g\n",
    //     DoCool_u_old_input, DoCool_rho_input, DoCool_dt_input, DoCool_ne_guess_input);

    //endrun(12);
    throw runtime_error("failed to converge in convert_u_to_temp()");
  }
  return temp;
}


pair<cooling::cooling_abundances, cooling::cooling_rates> 
cooling::ComputeAbundancesAndRates(double T, double rho, 
				   double ne_guess) const
{
  cooling_abundances a;

  if(ne_guess == 0)
    ne_guess = 1.0;
  a.ne = ne_guess;


  cooling_rates v = ComputeRates(T);


  a.nH0 = v.AlphaHp / (v.AlphaHp + v.GammaeH0 );	/* eqn (33) */
  a.nHp = 1.0 - a.nH0;		/* eqn (34) */

  if(v.GammaeHe0 <= SMALLNUM)	/* no ionization at all */
    {
      a.nHep = 0.0;
      a.nHepp = 0.0;
      a.nHe0 = yhelium;
    }
  else
    {
      a.nHep = yhelium / (1.0 + (v.AlphaHep + v.Alphad) / v.GammaeHe0 +
			  v.GammaeHep / v.AlphaHepp);	/* eqn (35) */
      a.nHe0 = a.nHep * (v.AlphaHep + v.Alphad) / v.GammaeHe0;	/* eqn (36) */
      a.nHepp = a.nHep * v.GammaeHep / v.AlphaHepp;	/* eqn (37) */
    }

  a.ne = a.nHp + a.nHep + 2 * a.nHepp;	/* eqn (38) */

  return make_pair(a, v);
}

cooling::cooling_rates cooling::ComputeRates(double T) const
{

#ifdef NEW_RATES
  double dE, P, A, X, K, U, T_eV;
  double b0, b1, b2, b3, b4, b5, c0, c1, c2, c3, c4, c5, y;	// used in Scholz-Walter fit 
  double E1s_2, Gamma1s_2s, Gamma1s_2p;
#endif

  cooling_rates v;

  const double Tfact = 1.0 / (1 + sqrt(T / 1.0e5));

  if(118348 / T < 70)
    v.BetaH0 = 7.5e-19 * exp(-118348 / T) * Tfact;

#ifdef NEW_RATES
  // Scholtz-Walters 91 fit
  if(T >= 2.0e3 && T < 1e8) {
    if(T >= 2.0e3 && T < 6.0e4) {
      b0 = -3.299613e1;
      b1 = 1.858848e1;
      b2 = -6.052265;
      b3 = 8.603783e-1;
      b4 = -5.717760e-2;
      b5 = 1.451330e-3;

      c0 = -1.630155e2;
      c1 = 8.795711e1;
      c2 = -2.057117e1;
      c3 = 2.359573;
      c4 = -1.339059e-1;
      c5 = 3.021507e-3;
    }
    else {
      if(T >= 6.0e4 && T < 6.0e6) {
	b0 = 2.869759e2;
	b1 = -1.077956e2;
	b2 = 1.524107e1;
	b3 = -1.080538;
	b4 = 3.836975e-2;
	b5 = -5.467273e-4;

	c0 = 5.279996e2;
	c1 = -1.939399e2;
	c2 = 2.718982e1;
	c3 = -1.883399;
	c4 = 6.462462e-2;
	c5 = -8.811076e-4;
      }
      else {
	b0 = -2.7604708e3;
	b1 = 7.9339351e2;
	b2 = -9.1198462e1;
	b3 = 5.1993362;
	b4 = -1.4685343e-1;
	b5 = 1.6404093e-3;

	c0 = -2.8133632e3;
	c1 = 8.1509685e2;
	c2 = -9.4418414e1;
	c3 = 5.4280565;
	c4 = -1.5467120e-1;
	c5 = 1.7439112e-3;
      }

      y = log(T);
      E1s_2 = 10.2;	/* eV */

      Gamma1s_2s =
	exp(b0 + b1 * y + b2 * y * y + b3 * y * y * y + b4 * y * y * y * y + b5 * y * y * y * y * y);
      Gamma1s_2p =
	exp(c0 + c1 * y + c2 * y * y + c3 * y * y * y + c4 * y * y * y * y + c5 * y * y * y * y * y);

      T_eV = T / eV_to_K;

      v.BetaH0 = E1s_2 * eV_to_erg * (Gamma1s_2s + Gamma1s_2p) * exp(-E1s_2 / T_eV);
    }
  }
#endif


  if(473638 / T < 70)
    v.BetaHep = 5.54e-17 * pow(T, -0.397) * exp(-473638 / T) * Tfact;

  v.Betaff = 1.43e-27 * sqrt(T) * (1.1 + 0.34 * exp(-(5.5 - log10(T)) * (5.5 - log10(T)) / 3));


#ifdef NEW_RATES
  v.AlphaHp = 6.28e-11 * pow(T / 1000, -0.2) / (1. + pow(T / 1.0e6, 0.7)) / sqrt(T);
#else
  v.AlphaHp = 8.4e-11 * pow(T / 1000, -0.2) / (1. + pow(T / 1.0e6, 0.7)) / sqrt(T);	// old Cen92 fit 
#endif


  v.AlphaHep = 1.5e-10 * pow(T, -0.6353);


#ifdef NEW_RATES
  v.AlphaHepp = 3.36e-10 * pow(T / 1000, -0.2) / (1. + pow(T / 4.0e6, 0.7)) / sqrt(T);
#else
  v.AlphaHepp = 4. * v.AlphaHp;	// old Cen92 fit
#endif

  if(470000 / T < 70)
    v.Alphad = 1.9e-3 * pow(T, -1.5) * exp(-470000 / T) * (1. + 0.3 * exp(-94000 / T));


#ifdef NEW_RATES
  T_eV = T / eV_to_K;

  // Voronov 97 fit
  // hydrogen 
  dE = 13.6;
  P = 0.0;
  A = 0.291e-7;
  X = 0.232;
  K = 0.39;

  U = dE / T_eV;
  v.GammaeH0 = A * (1.0 + P * sqrt(U)) * pow(U, K) * exp(-U) / (X + U);

  // Helium
  dE = 24.6;
  P = 0.0;
  A = 0.175e-7;
  X = 0.18;
  K = 0.35;

  U = dE / T_eV;
  v.GammaeHe0 = A * (1.0 + P * sqrt(U)) * pow(U, K) * exp(-U) / (X + U);

  // Helium II
  dE = 54.4;
  P = 1.0;
  A = 0.205e-8;
  X = 0.265;
  K = 0.25;

  U = dE / T_eV;
  v.GammaeHep = A * (1.0 + P * sqrt(U)) * pow(U, K) * exp(-U) / (X + U);

#else
  if(157809.1 / T < 70)
    v.GammaeH0 = 5.85e-11 * sqrt(T) * exp(-157809.1 / T) * Tfact;

  if(285335.4 / T < 70)
    v.GammaeHe0 = 2.38e-11 * sqrt(T) * exp(-285335.4 / T) * Tfact;

  if(631515.0 / T < 70)
    v.GammaeHep = 5.68e-12 * sqrt(T) * exp(-631515.0 / T) * Tfact;
#endif

  return v;
}




